package common.apkeep;

import jdd.bdd.BDD;

import java.util.*;

public class PPM {
    public HashMap<String, LinkedHashSet<Integer>> portToPreds;
    public HashMap<Integer, TreeSet<String>> predToPorts;
    public HashMap<TreeSet<String>, HashSet<Integer>> portsToPreds;
    private final BDD bdd;
    private LinkedList<Integer> transferredPreds;

    public PPM(BDD bdd) {
        this.portToPreds = new HashMap<>();
        this.predToPorts = new HashMap<>();
        this.portsToPreds = new HashMap<>();
        this.bdd = bdd;
    }

    public void update(int delta, String from, String to, LinkedList<Integer> transferredPreds) {
        this.transferredPreds = transferredPreds;
        HashSet<Integer> oldPortPreds = new HashSet<>(portToPreds.get(from));
        for (int p : oldPortPreds) {
            int intersection = bdd.ref(bdd.and(p, delta));
            if (intersection == 0) {
                bdd.deref(intersection);
                continue;
            }

            if (intersection != p) {
                int tmp = bdd.ref(bdd.not(delta));
                int tmp1 = bdd.ref(bdd.and(p, tmp));
                bdd.deref(tmp);
                this.split(p, intersection, tmp1);
            }

            this.transfer(intersection, from, to);

            HashSet<Integer> newPredicates = new HashSet<>(this.portToPreds.get(to));
            int oldTarget = intersection;
            for (int t : newPredicates) {
                if (oldTarget != t && predToPorts.get(intersection).equals(predToPorts.get(t))) {
                        intersection = this.merge(intersection, t, bdd.ref(bdd.or(intersection, t)));
                }
            }
//            this.checkMerge();
        }
    }

    private void putPredToPorts(int pred, TreeSet<String> ports) {
        if (portsToPreds.containsKey(ports)) {
            portsToPreds.get(ports).remove(pred);
        }
        predToPorts.put(pred, ports);
        portsToPreds.computeIfAbsent(predToPorts.get(pred), k -> new HashSet<>());
        portsToPreds.get(ports).add(pred);
    }

    private void removePredToPorts(int pred) {
        portsToPreds.get(predToPorts.get(pred)).remove(pred);
        predToPorts.remove(pred);
    }

    private void addPredToPorts(int p, String port) {
        portsToPreds.get(predToPorts.get(p)).remove(p);
        predToPorts.get(p).add(port);
        putPredToPorts(p, predToPorts.get(p));
    }

    private void subPredToPorts(int p, String port) {
        portsToPreds.get(predToPorts.get(p)).remove(p);
        predToPorts.get(p).remove(port);
        putPredToPorts(p, predToPorts.get(p));
    }

    private void split(int p, int p1, int p2) {
        for (String port : this.predToPorts.get(p)) {
            this.portToPreds.get(port).addAll(Arrays.asList(p1, p2));
            this.portToPreds.get(port).remove(p);
        }
//        putPredToPorts(p1, this.predToPorts.get(p));
//        putPredToPorts(p2, this.predToPorts.get(p));
//        removePredToPorts(p);
        this.predToPorts.put(p1, new TreeSet<> (this.predToPorts.get(p)));
        this.predToPorts.put(p2, new TreeSet<> (this.predToPorts.get(p)));
        this.predToPorts.remove(p);

//        this.portsToPreds.get(this.predToPorts.get(p1)).remove(p);
//        this.portsToPreds.get(this.predToPorts.get(p1)).addAll(Arrays.asList(p1, p2));

        // here we derefer p
        bdd.deref(p);

        if (this.transferredPreds.contains(p)) {
            this.transferredPreds.addAll(Arrays.asList(p1, p2));
            this.transferredPreds.remove(p);
        }
    }

    private void transfer(int p, String from, String to) {
        this.portToPreds.get(from).remove(p);
        this.portToPreds.get(to).add(p);

//        if (this.portsToPreds.containsKey(this.predToPorts.get(p))) {
        assert this.predToPorts.get(p) != null;
        if (this.portsToPreds.get(this.predToPorts.get(p)) == null) {
//            System.out.println("");
        }
//        removePredToPorts(p);
//            this.portsToPreds.get(this.predToPorts.get(p)).remove(p);
//        }
//        addPredToPorts(p, to);
//        subPredToPorts(p, from);

        this.predToPorts.get(p).add(to);
        this.predToPorts.get(p).remove(from);
//        portsToPreds.computeIfAbsent(predToPorts.get(p), k -> new HashSet<>());
//        this.portsToPreds.get(this.predToPorts.get(p)).add(p);

        this.transferredPreds.add(p);
    }

    private void checkMerge() {
        HashMap<TreeSet<String>, LinkedList<Integer>> ptp = new HashMap<>();
        for (Map.Entry<Integer, TreeSet<String>> entry : this.predToPorts.entrySet()) {
            ptp.computeIfAbsent(entry.getValue(), k -> new LinkedList<>());
            ptp.get(entry.getValue()).add(entry.getKey());
        }
//        HashMap<TreeSet<Port>, HashSet<Integer>> pp = (HashMap<TreeSet<Port>, HashSet<Integer>>) portsToPreds.clone();
        for (LinkedList<Integer> preds : ptp.values()) {
            if (preds.size() > 5) {
                this.mergeAll(preds);
            }
        }
    }

    public void finishMerge() {
        for (Map.Entry<TreeSet<String>, HashSet<Integer>> entry : this.portsToPreds.entrySet()) {
            this.mergeAll(new LinkedList<>(entry.getValue()));
        }
        this.portsToPreds.clear();
    }

    private int mergeAll(LinkedList<Integer> preds) {
        int mreged = preds.stream().reduce(0, (p1, p2) -> bdd.ref(bdd.or(p1, p2)));
        for (String port : this.predToPorts.get(preds.get(0))) {
            this.portToPreds.get(port).add(mreged);
            this.portToPreds.get(port).removeAll(preds);
        }
        this.predToPorts.put(mreged, (TreeSet<String>) this.predToPorts.get(preds.get(0)).clone());
        for (int p : preds) {
            this.predToPorts.remove(p);
        }
        return mreged;
//        HashSet<Integer> cpreds = (HashSet<Integer>) preds.clone();
//        return cpreds.stream().reduce(0, (p1, p2) -> this.merge(p1, p2, bdd.ref(bdd.or(p1, p2))));
    }

    private int merge(int p1, int p2, int p) {
        if (p1 == 0) {
            return p2;
        }
        for (String port : this.predToPorts.get(p1)) {
            this.portToPreds.get(port).add(p);
            this.portToPreds.get(port).removeAll(Arrays.asList(p1, p2));
        }
//        putPredToPorts(p, this.predToPorts.get(p1));
//        removePredToPorts(p1);
//        removePredToPorts(p2);
        this.predToPorts.put(p, new TreeSet<>(this.predToPorts.get(p1)));
        this.predToPorts.remove(p1);
        this.predToPorts.remove(p2);

//        this.portsToPreds.get(this.predToPorts.get(p)).add(p);
//        this.portsToPreds.get(this.predToPorts.get(p)).removeAll(Arrays.asList(p1, p2));

        if (this.transferredPreds.contains(p1) || this.transferredPreds.contains(p2)) {
            this.transferredPreds.add(p);
            this.transferredPreds.removeAll(Arrays.asList(p1, p2));
        }
        return p;
    }
}
